from typing import Callable

import numpy as np
from numpy import linalg as LA
import math
from collections import deque
from scipy.spatial.transform import Rotation as R
from sklearn.metrics import pairwise_distances
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

from dora import DoraStatus
from spline_utils import calc_spline_course

# Input type
USE_TRAJECTORY = True

# Throttle PID Control
LON_CONTROL_PID_P = 0.70
LON_CONTROL_PID_I = 0.05
LON_CONTROL_PID_D = 0.05
LON_CONTROL_MAX_THROTTLE = 1.0
LON_CONTROL_MAX_BRAKE = 1.0
LON_CONTROL_SPEED_LOOKAHEAD_NUM = 1
LON_CONTROL_ENABLE_START_ASSIST = True
LON_CONTROL_START_ASSIST_THROTTLE = 0.45
LON_CONTROL_START_ASSIST_SPEED = 2.0 # m/s

# Vehicle
VECHILE_DT = 1.0 / 20  # s
VECHILE_DIST_THRESHOLD = 0.1  # m
VECHILE_FRONT_DIST_TO_CENTER = 1.4775 # m
VECHILE_MAX_STEER = 1.0  # rad
VECHILE_MAX_SPEED = 10  # m/s
VECHILE_MAX_ACCEL = 10  # m/s^2

# Stanley Control
STANLEY_HEADING_GAIN = 0.8
STANLEY_CONTROL_GAIN = 0.8

# Lateral lookahead
LAT_LOOKAHEAD_DIST = 4.0 # m

# Debug
DEBUG = False


class SecondOrderLowPass:
    """
    2nd-order Low Pass Filter
    """

    def __init__(self, cutoff_freq, sample_freq):
        self.has_initialized = False
        self.configure(cutoff_freq, sample_freq)
        self.i = np.empty((2,))
        self.o = np.empty((2,))

    def configure(self, cutoff_freq, sample_freq):
        assert cutoff_freq > 0.0, "Cuffoff frequency should be positive!"
        assert sample_freq > 0.0, "Sampling frequency should be positive!"

        self.a = np.empty((2,))
        self.b = np.empty((2,))
        K = np.tan(math.pi * cutoff_freq / sample_freq)
        Q = 1.0 / (2.0 * np.cos(math.pi / 4))
        poly = K * K + K / Q + 1.0
        self.a[0] = 2.0 * (K * K - 1.0) / poly
        self.a[1] = (K * K - K / Q + 1.0) / poly
        self.b[0] = K * K / poly
        self.b[1] = 2.0 * self.b[0]

    def apply(self, sample):
        if not self.has_initialized:
            self.has_initialized = True
            self.i[0], self.i[1] = sample, sample
            self.o[0], self.o[1] = sample, sample
            return sample

        out = (
            self.b[0] * sample
            + self.b[1] * self.i[0]
            + self.b[0] * self.i[1]
            - self.a[0] * self.o[0]
            - self.a[1] * self.o[1]
        )

        self.i[1] = self.i[0]
        self.i[0] = sample
        self.o[1] = self.o[0]
        self.o[0] = out
        return out

    def get(self):
        return self.o[0]

    def reset(self, value=None):
        if value is None:
            self.i = np.empty((2,))
            self.o = np.empty((2,))
            self.has_initialized = False
        else:
            self.i[0], self.i[1] = value, value
            self.o[0], self.o[1] = value, value
            self.has_initialized = True


class PID:
    """
    PID Control Implementation
    """

    def __init__(self, K_P=1.0, K_I=0.0, K_D=0.0, dt=0.05, output_limit=1.0):
        self.k_p = K_P
        self.k_i = K_I
        self.k_d = K_D
        self.dt = dt
        self.output_limit = output_limit

        # back-calculation anti-windup
        self.k_b = 1 / self.k_p
        if math.isnan(self.k_b):
            self.k_b = 0.0

        self.reset()

    def reset(self):
        self.last_error = 0.0
        self.error_sum = 0.0
        self.back_calculation = 0.0
        self.lpf = SecondOrderLowPass(cutoff_freq=5, sample_freq=1 / self.dt)

    def update(self, target, current):
        error = target - current

        # Proportional term
        output_P = self.k_p * error

        # Integral term
        self.error_sum += error * self.dt
        self.error_sum += self.back_calculation * self.k_b * self.dt
        output_I = self.k_i * self.error_sum

        # Derivative term
        d_error = (error - self.last_error) / self.dt
        d_error = self.lpf.apply(d_error)
        self.last_error = error
        output_D = self.k_d * d_error

        # output
        output = output_P + output_I + output_D

        # anti-windup using back-calculation
        output_saturated = np.clip(output, -self.output_limit, self.output_limit)
        self.update_saturation(output_saturated, output)

        if DEBUG:
            print(
                f"[Throttle PID Control] P: {output_P}, I: {output_I}, D: {output_D}, output: {output_saturated}"
            )

        return output_saturated

    def update_saturation(self, output_saturated, output):
        self.back_calculation = output_saturated - output

    def set_I_term(self, value):
        if self.k_i != 0:
            self.error_sum = value / self.k_i

        return True

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x)) 

def compute_throttle_and_brake(pid: PID, current_speed: float, target_speed: float, end_speed: float):
    """Computes the throttle/brake required to reach the target speed."""
    if current_speed < 0:
        if DEBUG:
            print(f"[Throttle PID Control] Current speed is negative: {current_speed}")
        non_negative_speed = 0.0
    else:
        non_negative_speed = current_speed

    # if target_speed == 0.0:
    #     pid.reset()

    # Check for start assist condition
    if (
        LON_CONTROL_ENABLE_START_ASSIST
        # and non_negative_speed < target_speed
        and non_negative_speed < end_speed
        and non_negative_speed < LON_CONTROL_START_ASSIST_SPEED
    ):
        if DEBUG:
            print(f"[Throttle PID Control] START_ASSIST is active")
        speed_ratio = non_negative_speed / LON_CONTROL_START_ASSIST_SPEED
        throttle_gain = sigmoid(speed_ratio * 10)
        throttle = LON_CONTROL_START_ASSIST_THROTTLE * throttle_gain
        brake = 0.0
        pid.reset()
        pid.set_I_term(throttle)
    else:
        acceleration = pid.update(target_speed, non_negative_speed)
        if acceleration >= 0.0:
            throttle = min(acceleration, LON_CONTROL_MAX_THROTTLE)
            brake = 0.0
        else:
            throttle = 0.0
            brake = min(abs(acceleration), LON_CONTROL_MAX_BRAKE)

    # Keep the brake pressed when stopped or when sliding back on a hill.
    if (current_speed < 1 and target_speed == 0) or current_speed < -0.3:
        brake = 1.0
    return throttle, brake


def wrap_angle(angle):
    """Wrap angle to -pi to pi."""
    return (angle + np.pi) % (2 * np.pi) - np.pi


def get_angle_difference(angle1, angle2):
    """Get angle difference."""
    vector1 = np.array([math.cos(angle1), math.sin(angle1)])
    vector2 = np.array([math.cos(angle2), math.sin(angle2)])
    dot_product = np.dot(vector1, vector2)
    sin_angle_diff_sign = np.sign(math.sin(angle1 - angle2))
    dot_product = min(dot_product, 1)
    dot_product = max(dot_product, -1)
    angle_v1_to_v2 = math.acos(dot_product) * sin_angle_diff_sign

    return angle_v1_to_v2


def get_projection_point(x_m, y_m, yaw_m, k_m, x, y):
    """Get projection point on the curve."""
    d_vector = np.array([x - x_m, y - y_m])
    tau_vector = np.array([math.cos(yaw_m), math.sin(yaw_m)])
    p_proj = np.array([x_m, y_m]) + (d_vector.dot(tau_vector)) * tau_vector
    yaw_proj = yaw_m + k_m * (d_vector.dot(tau_vector))

    return p_proj, yaw_proj


class Operator:
    """
    Compute the throttle, target angle and brake given a `position`, a `speed` and a `waypoints`.
    """

    def __init__(self):
        self.position = np.array([])
        self.speed = np.array([])
        self.last_target_speed = 0.0
        self.last_target_angle = 0.0

        # Longitudinal PID Controller
        self.longitudinal_pid_controller = PID(
            LON_CONTROL_PID_P,
            LON_CONTROL_PID_I,
            LON_CONTROL_PID_D,
            VECHILE_DT,
            output_limit=1.0,
        )

        ## DEBUG
        if DEBUG:
            max_len = 150
            self.speed_cmd_history = deque(maxlen=max_len)
            self.actual_speed_history = deque(maxlen=max_len)

            self.fig, self.axis = plt.subplots(1, 1)
            self.fig.canvas.set_window_title("PID Monitor")
            self.fig.canvas.manager.window.wm_geometry("+100+800")
            (self.speed_cmd_handle,) = self.axis.plot(
                self.speed_cmd_history, "b--", linewidth=1.0, label="target"
            )
            (self.actual_speed_handle,) = self.axis.plot(
                self.actual_speed_history, "r", linewidth=1.0, label="actual"
            )
            self.axis.legend()

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):
        """Handle input.
        Args:
            dora_input["id"](str): Id of the input declared in the yaml configuration
            dora_input["data"] (bytes): Bytes message of the input
            send_output (Callable[[str, bytes]]): Function enabling sending output back to dora.
        """

        if dora_input["id"] == "position":
            self.position = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "speed":
            self.speed = np.frombuffer(dora_input["data"], np.float32)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "waypoints":
            if USE_TRAJECTORY:
                return DoraStatus.CONTINUE
            path = np.frombuffer(dora_input["data"], np.float32).reshape((-1, 3))

        elif dora_input["id"] == "trajectory":
            if not USE_TRAJECTORY:
                return DoraStatus.CONTINUE
            trajectory = np.frombuffer(dora_input["data"], np.float32).reshape((-1, 7))

        # If NO localization
        if len(self.position) == 0 or len(self.speed) == 0:
            self.last_target_speed = 0.0
            self.last_target_angle = 0.0
            self.longitudinal_pid_controller.reset()

            send_output(
                "control",
                np.array([0, 0, 1], np.float16).tobytes(),
                {},
            )
            return DoraStatus.CONTINUE

        # If NO waypoints or trajectory
        if (USE_TRAJECTORY and len(trajectory) == 0) or (
            not USE_TRAJECTORY and len(path) == 0
        ):
            self.last_target_speed = 0.0
            self.last_target_angle = 0.0
            self.longitudinal_pid_controller.reset()

            send_output(
                "control",
                np.array([0, 0, 1], np.float16).tobytes(),
                {},
            )
            return DoraStatus.CONTINUE

        # Get ego position and front axle position
        [x, y, _, qx, qy, qz, qw] = self.position
        [_, _, yaw] = R.from_quat([qx, qy, qz, qw]).as_euler("xyz", degrees=False)
        # positin of front axle
        front_x = x + VECHILE_FRONT_DIST_TO_CENTER * np.cos(yaw)
        front_y = y + VECHILE_FRONT_DIST_TO_CENTER * np.sin(yaw)

        # Extract points and target speeds
        if USE_TRAJECTORY:
            # trajectory = [x, y, vel, yaw, k, s, t]
            waypoints = trajectory[:, :2]
            target_speeds = trajectory[:, 2]
        else:
            # path = [x, y, vel]
            waypoints = path[:, :2]
            target_speeds = path[:, 2]

        # Steering control
        if (
            len(waypoints) == 1
        ):  # If waypoints only contain 1 point, use pure-pursuit tracking algorithm
            dist_to_end = np.sqrt(
                (x - waypoints[0, 0]) ** 2 + (y - waypoints[0, 1]) ** 2
            )

            # compute the angle of steering
            target_vector = waypoints[0] - np.array([x, y])
            target_angle = math.atan2(target_vector[1], target_vector[0]) - yaw
            target_angle = np.clip(
                wrap_angle(target_angle), -VECHILE_MAX_STEER, VECHILE_MAX_STEER
            )

            # target speed
            target_speed = max(1.0, target_speeds[0])
            end_speed = target_speed

        else:  # Otherwise, we use stanley tracking algorithm
            # tracked_waypoints = [x, y, yaw, k, vel]
            if USE_TRAJECTORY:
                tracked_waypoints = trajectory[:, [0, 1, 3, 4, 2]]
            else:
                rx, ry, ryaw, rk, rs, sp = calc_spline_course(
                    waypoints[:, 0], waypoints[:, 1], ds=0.1
                )
                rs = np.clip(np.asarray(rs), sp.s[0], sp.s[-1])
                if len(sp.singular_idx) > 0:
                    target_speeds = np.delete(target_speeds, sp.singular_idx, axis=0)
                speed_fit_fcn = interp1d(sp.s, target_speeds)
                resampled_speeds = speed_fit_fcn(rs)
                tracked_waypoints = np.array([rx, ry, ryaw, rk, resampled_speeds]).T

            # get tracking status
            center_distances = pairwise_distances(
                tracked_waypoints[:, :2], np.array([[x, y]])
            ).T[0]
            front_distances = pairwise_distances(
                tracked_waypoints[:, :2],
                np.array([[front_x, front_y]]),
            ).T[0]
            center_pt_idx = np.argmin(center_distances)
            front_pt_idx = np.argmin(front_distances)
            dist_to_end = center_distances[-1]

            # Get target speed
            target_speed = tracked_waypoints[
                min(
                    center_pt_idx + LON_CONTROL_SPEED_LOOKAHEAD_NUM,
                    len(tracked_waypoints) - 1,
                ),
                4,
            ]
            end_speed = tracked_waypoints[-1, 4]
            
            # Take average of target speed
            # target_speed = (target_speed + tracked_waypoints[-1, 4]) / 2
            # target_speed = tracked_waypoints[-1, 4]

            # if target_speed is close to zero
            if target_speed < 1e-3:
                target_speed = 0.0

            # # matching point
            # front_axle_path_pose = tracked_waypoints[front_pt_idx]
            # x_m, y_m, yaw_m, k_m = (
            #     front_axle_path_pose[0],
            #     front_axle_path_pose[1],
            #     front_axle_path_pose[2],
            #     front_axle_path_pose[3],
            # )

            # # get projection point
            # p_proj, yaw_proj = get_projection_point(
            #     x_m, y_m, yaw_m, k_m, front_x, front_y
            # )
            # front_axle_vec = np.array(
            #     [-np.cos(yaw + np.pi / 2), -np.sin(yaw + np.pi / 2)]
            # )
            # d_vector = np.array([front_x, front_y]) - p_proj
            # error_front_axle = np.dot(d_vector, front_axle_vec)

            # # theta_e corrects the heading error
            # target_yaw = wrap_angle(yaw_proj)
            # theta_e = get_angle_difference(target_yaw, yaw) * STANLEY_HEADING_GAIN

            # # theta_d corrects the cross track error
            # theta_d = np.arctan2(
            #     STANLEY_CONTROL_GAIN * error_front_axle, max(0.3, self.speed[0])
            # )
            # theta_d = wrap_angle(theta_d)

            # # steering control
            # target_angle = wrap_angle(theta_e + theta_d)
            # target_angle = np.clip(target_angle, -VECHILE_MAX_STEER, VECHILE_MAX_STEER)

            # extend the tracked waypoints
            extend_tracked_waypoints = tracked_waypoints.copy()
            extend_dir = tracked_waypoints[-1, :2] - tracked_waypoints[-2, :2]
            extend_dir = extend_dir / np.linalg.norm(extend_dir)
            extend_len = LAT_LOOKAHEAD_DIST
            step_size = 0.2
            num_new_pts = math.ceil(extend_len / step_size)
            for _ in range(num_new_pts):
                new_xy = extend_tracked_waypoints[-1, :2] + extend_dir * step_size
                extend_tracked_waypoints = np.vstack(
                    [
                        extend_tracked_waypoints, 
                        np.array(
                            [
                                new_xy[0],
                                new_xy[1],
                                extend_tracked_waypoints[-1, 2],
                                0.0,
                                extend_tracked_waypoints[-1, 4],
                            ]
                        )
                    ]
                )

            # use pure-pursuit algorithm
            lookahead_idx = center_pt_idx
            distance = 0.0
            while lookahead_idx < len(extend_tracked_waypoints) - 1:
                distance += LA.norm(
                    extend_tracked_waypoints[lookahead_idx + 1, :2] - extend_tracked_waypoints[lookahead_idx, :2]
                )
                lookahead_idx += 1
                if distance >= LAT_LOOKAHEAD_DIST:
                    break

            target_vector = extend_tracked_waypoints[lookahead_idx, :2] - np.array([x, y])
            target_angle = math.atan2(target_vector[1], target_vector[0]) - yaw
            target_angle = np.clip(
                wrap_angle(target_angle), -VECHILE_MAX_STEER, VECHILE_MAX_STEER
            )

        # check if tracking done
        if dist_to_end < VECHILE_DIST_THRESHOLD:
            target_speed = 0.0

        # Speed control
        # rate limiter
        if target_speed != 0.0:  # except stop
            max_speed_increment = abs(VECHILE_MAX_ACCEL * VECHILE_DT)
            target_speed = self.last_target_speed + np.clip(
                target_speed - self.last_target_speed,
                -max_speed_increment,
                max_speed_increment,
            )

        # Compute throttle and brake
        throttle, brake = compute_throttle_and_brake(
            self.longitudinal_pid_controller,
            self.speed[0],
            np.clip(target_speed, -VECHILE_MAX_SPEED, VECHILE_MAX_SPEED),
            end_speed,
        )

        self.last_target_speed = target_speed
        self.last_target_angle = target_angle

        # Send output
        control_cmd = np.array([throttle, target_angle, brake], np.float16)
        send_output(
            "control",
            control_cmd.tobytes(),
            {},
        )

        ## DEBUG
        if DEBUG:
            self.speed_cmd_history.append(target_speed)
            self.speed_cmd_handle.set_xdata(range(len(self.speed_cmd_history)))
            self.speed_cmd_handle.set_ydata(self.speed_cmd_history)

            self.actual_speed_history.append(self.speed[0])
            self.actual_speed_handle.set_xdata(range(len(self.actual_speed_history)))
            self.actual_speed_handle.set_ydata(self.actual_speed_history)

            self.axis.set_title(f"throttle={throttle:.2f}, brake={brake:.2f}")
            self.axis.set_ylabel("speed [m/s]")
            self.axis.relim()
            self.axis.autoscale_view()
            plt.pause(1e-5)

        return DoraStatus.CONTINUE
